## master preamble ------------------------------------------------------
rm(list=ls(all=TRUE)) 

## inherit directories from stata -----------
args = commandArgs(trailingOnly = "TRUE")
if (length(args)) {
  dir.proj = paste0(args[1],'/')
  dir.lib  = args[2]
} 

## if loading from library -------------------
if(dir.lib!='') {
  ## load depenencies --------------------
  library(R.utils,lib.loc=dir.lib)
  library(ggplot2,lib.loc=dir.lib)
  library(tidyr,lib.loc=dir.lib)
  library(stringr,lib.loc=dir.lib)
  library(forcats,lib.loc=dir.lib)
  library(lubridate,lib.loc=dir.lib)
  library(haven,lib.loc=dir.lib)
  ## load main packages ------------------
  library(rio,lib.loc=dir.lib)
  library(dplyr,lib.loc=dir.lib)
  library(tidyverse,lib.loc=dir.lib)
  library(fixest,lib.loc=dir.lib)
## otherwise: simple load -----------------
} else {
  library(rio)
  library(dplyr)
  library(tidyverse)
  library(fixest)
}
## --------------------------------------------------------------------


## file preamble ------------------------------------------------------
## set randomization seed --------------
set.seed(661)
## directories --------
dir.data = paste0(dir.proj,'data/')
dir.est  = paste0(dir.proj,'estimates/')
dir.out  = paste0(dir.proj,'output/')
dir.temp = paste0(dir.proj,'_temp/')
## varlist -----------
cov.list = c('female','age','agesq','age_miss',
             'race_b','race_h','race_o','race_u','priorprison','local',
             'logzipincome','zipincome_miss','logprice','veh_miss',
             'speed_py1','other_py1','crash_py1')
## --------------------------------------------------------------------



### Setup data -----------------------------------------------------------------
main <- rio::import(paste0(dir.data,'out/4-main.dta'))
main <- mutate(
  main,
  D = as.numeric(discount==1))
main <- dplyr::select(
  main,
  D,all_of(cov.list),totfe,officerid)


######################
### Fit 1[bunch] to officer FE, covs, beat-shift FE ------------
fit.all = feols(
  as.formula(paste('D ~ 0 + i(officerid) +',paste(cov.list,collapse='+'),'|totfe')),main)

######################
### Joint test of officer FE -----------------------
fit.test = wald(fit.all,drop=cov.list,print=F)

######################
### Store fit info ---------------------------------
r2.tot = as.numeric(r2(fit.all,type='r2'))
r2.fe  = r2.tot - as.numeric(r2(fit.all,type='wr2'))

## Redo without officer FE to get covs R2 -----
fit.drop = feols(
  as.formula(paste('D ~ ',paste(cov.list,collapse='+'),'|totfe')),
  main,vcov=~totfe)
r2.cov = as.numeric(r2(fit.drop,type='wr2'))

## Other way of checking cov R2? ---------------
fit.drop = feols(
  as.formula(paste('D ~ 0 | totfe+officerid')),main)
r2.covalt = r2.tot - as.numeric(r2(fit.drop,type='r2'))

## Officer fe is remainder ----------
r2.off = r2.tot-(r2.fe+r2.cov)

## Export FE estimates --------------
out.file <- file(paste0(dir.out,'apx_iv/rsquared.txt'))
writeLines(c(
  paste0('Total R2      = ',round(r2.tot,4)),
  paste0('Beat-Shift R2 = ',round(r2.fe, 4)),
  paste0('Driver R2     = ',round(r2.cov,4)),
  paste0('Officer R2    = ',round(r2.off,4)),
  paste0('Officer Share = ',round(r2.off/r2.tot,4)),
  paste0('Joint Test: F = ',round(fit.test$stat,4),' (p = ',fit.test$p,')')),
  out.file)
close(out.file)
### -------------------------------------------------
######################


###########################
### Store theta estimates ---------------------------
coef = filter(
  data.frame(
  officerid = as.numeric(str_replace(names(fit.all$coefficients),'officerid::','')),
  est = fit.all$coeftable[,1],
  se  = fit.all$coeftable[,2],row.names=NULL),
  !is.na(officerid))
export(coef,paste0(dir.est,'getfe.dta'))
## ----------------------------------------------------



